package com.example.agent_springai.contoller;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@RequiredArgsConstructor
@RestController
@RequestMapping("/ai")
public class KnowledgeDBInsertController {

    @Autowired
    private VectorStore vectorStore;

    @GetMapping("/pdf")
    public String testRedisVectorStore(@RequestParam String prompt){
        initKnowledgeDB();

        SearchRequest request = SearchRequest.builder()
                .query(prompt) // 查询
                .topK(5)  // 返回的相似文档数量
                .similarityThreshold(0.6) // 相似度阈值
                //.filterExpression("file_name == '中二知识笔记.pdf'")  // 过滤条件
                .build();
        List<Document> docs = vectorStore.similaritySearch(request);
        if (docs == null) {
            return "没有搜索到任何内容";
        }
        for (Document doc : docs) {
            System.out.println(doc.getId());
            System.out.println(doc.getScore());
            System.out.println(doc.getText());
        }
        return docs.get(0).getText();
    }

    public void initKnowledgeDB() {
        Resource resource = new ClassPathResource("knowledgeDB/S120_故障手册.pdf");
        //创建pdf读取器
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource,
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1)//每一页为一个片段
                        .build()
        );
        //拆分pdf
        List<Document> documents = reader.read();
        //写入向量数据库
        vectorStore.add(documents);
    }
}
